import pandas as pd
from pandas import DataFrame
from app_config import get_engine_ts, get_pro
from pyecharts.charts import Line
from pyecharts import options as opts
from pyecharts.options import LabelOpts
from datetime import datetime
from dateutil.relativedelta import relativedelta
import os

engine = get_engine_ts()
pro = get_pro()


# IC2508.CFX 500

def get_data(path: str) -> DataFrame:
    # 计算每个月差值 生成图表
    df_data: DataFrame = pd.read_excel(path)
    df_data = df_data[['日期Date', '收盘Close']]
    df_data.rename(columns={
        '日期Date': 'date',
        '收盘Close': 'close'
    }, inplace=True)
    return df_data


def calculate_new_column(ts_code, date, dist_code_data: dict):
    if ts_code is None:
        return None
    code_dataframe = dist_code_data.get(ts_code)

    if code_dataframe is None:
        return None

    # 找到data列等于'202501'的行
    row = code_dataframe[code_dataframe['trade_date'] == date]
    # 如果找到行，返回close列的值
    if not row.empty:
        return row['close'].values[0]
    else:
        return None


def calculate_fut_code(date_str, product):
    date_obj = datetime.strptime(date_str[:6], "%Y%m")
    new_date_obj = date_obj + relativedelta(months=1)
    return product + new_date_obj.strftime("%y%m") + '.CFX'


def get_or_create_dataframe(product) -> DataFrame:
    if product == 'IC':
        file_path = 'fetch/000905.xlsx'
        t_code = '000905.SH'
    else:
        raise Exception
    if os.path.exists(file_path):
        df_excel = pd.read_excel(file_path)
    else:
        df_excel = get_pro().index_daily(ts_code=t_code)
        df_excel.to_excel(file_path, index=False)
    re = df_excel[['trade_date', 'close']]
    re.rename(columns={
        'trade_date': 'date',
        'close': 'close'
    }, inplace=True)
    return re


# 000905.SH
if __name__ == '__main__':
    prod = 'IC'
    df = get_or_create_dataframe(prod)
    df['date'] = df['date'].astype(str)
    df['fut_ts_code'] = df.apply(lambda row: calculate_fut_code(row['date'], prod), axis=1)

    sql_daily = f"""
       select ts_code,trade_date,close,settle from fut_daily where ts_code like '{prod}%'
       """
    
    df_daily = pd.read_sql_query(sql_daily, engine)

    grouped_dict = {key: group for key, group in df_daily.groupby('ts_code')}
    # 添加新列
    df['fut_close'] = df.apply(lambda row: calculate_new_column(row['fut_ts_code'], row['date'], grouped_dict), axis=1)
    df['delta'] = df['close'] - df['fut_close']
    df['delta_percent'] = df['delta'] / df['fut_close'] * 100

    line = Line()
    # line.set_series_opts(label_opts=opts.LabelOpts(is_show=False))
    line.set_global_opts(title_opts=opts.TitleOpts(title='工具栏显示')
                         , datazoom_opts=opts.DataZoomOpts(is_show=True,
                                                           range_start=0,
                                                           range_end=80,
                                                           orient='horizontal'
                                                           )
                         )
    line.add_xaxis(df['date'].tolist())
    line.add_yaxis(series_name="index",
                   y_axis=((df['close'] - 6000) / 1000).tolist(),
                   label_opts=LabelOpts(is_show=False),
                   tooltip_opts={"show": True})
    line.add_yaxis(series_name="fut",
                   y_axis=((df['fut_close'] - 6000) / 1000).tolist(),
                   label_opts=LabelOpts(is_show=False),
                   tooltip_opts={"show": True})
    line.add_yaxis(series_name="percent",
                   y_axis=df['delta_percent'].tolist(),
                   label_opts=LabelOpts(is_show=False),
                   tooltip_opts={"show": True})
    line.render("./f_chart/future_water2.html")
    print(df)
